import json
import logging
import time
from typing import Dict, Any, List, Union

import cv2
import numpy as np
from paddleocr import PaddleOCR

from app.utils.exceptions import ImageProcessingError
from app.utils.logger import get_logger
from app.utils.timing import timing_decorator

# 模块级logger，用于类内部使用
logger = get_logger(__name__)

# 用于缓存PaddleOCR实例的字典
ocr_cache = {}


class TextRecognitionModule:
    def __init__(self):
        pass

    # @timing_decorator(include_caller=True, log_level=logging.DEBUG)
    def recognize(
            self,
            image_data: Union[bytes, List[bytes]],
            language: str = "ch",
            detail: bool = False
    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
        """执行文本识别，支持单个或多个图像"""
        if isinstance(image_data, list):
            # 处理多个图像
            results = []
            for i, single_image_data in enumerate(image_data):
                logger.info(f"处理第 {i + 1} 个图像")
                result = self._recognize_single(single_image_data, language, detail)
                results.append(result)
            return results
        else:
            # 处理单个图像
            return self._recognize_single(image_data, language, detail)

    @timing_decorator(include_caller=True, log_level=logging.INFO, prefix="OCR处理")
    def _recognize_single(
            self,
            image_data: bytes,
            language: str,
            detail: bool
    ) -> Dict[str, Any]:
        """处理单个图像的文本识别"""
        try:
            # 将字节数据转换为OpenCV图像
            nparr = np.frombuffer(image_data, np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

            logger.info(f"图像尺寸: {img.shape if img is not None else 'None'}")
            if img is None:
                raise ImageProcessingError("无法解码上传的图像文件")

            # 获取或创建特定语言的OCR模型（带缓存机制）
            ocr_model = self._get_ocr_model(language)

            @timing_decorator(include_caller=True, log_level=logging.INFO ,prefix="推理分析")
            def predict(img, ocr_model):
                result = ocr_model.predict(
                    input=img,
                    use_textline_orientation=False,  # 禁用文本方向检测
                    use_doc_orientation_classify=False,  # 禁用文档预处理
                    use_doc_unwarping=False,  # 禁用文档矫正
                    return_word_box=True  # 返回单词/字符坐标
                )
                return result

            result = predict(img, ocr_model)
            # 输出推理参数
            # logger.info(f"推理参数: {ocr_model}")
            # logger.info("文本识别完成")

            return format_ocr_results_as_json(result)
        except Exception as e:
            logger.error(f"文本识别失败: {str(e)}")
            raise ImageProcessingError(f"Text recognition failed: {str(e)}")

    @timing_decorator(include_caller=True, log_level=logging.INFO, prefix="获取模型")
    def _get_ocr_model(self, language: str) -> PaddleOCR:
        """获取OCR模型实例（带缓存机制）"""
        global ocr_cache

        # 如果该语言的模型已存在缓存中，则直接返回
        if language in ocr_cache:
            # logger.info(f"使用缓存中的模型= {ocr_cache[language]}")
            return ocr_cache[language]

        # 否则创建新的模型实例并缓存
        # logger.info(f"准备使用{language}语言模型进行文本识别")
        start_time = time.time()
        ocr_model = PaddleOCR(
            text_detection_model_name="PP-OCRv5_server_det",
            text_recognition_model_name="PP-OCRv5_server_rec",
        )
        init_time = time.time() - start_time

        # 缓存模型实例和初始化时间
        ocr_cache[language] = ocr_model
        logger.info(f"模型初始化完成，Time: {init_time:.2f}s")
        return ocr_model


def format_ocr_results_as_json(results):
    """
    将PaddleOCR的识别结果格式化为JSON数据
    """
    data = []
    # 按指定格式处理结果
    for result in results:
        res_data = result.json['res']

        # 处理每个检测到的文本框
        dt_polys = res_data.get('dt_polys', [])
        rec_texts = res_data.get('rec_texts', [])
        rec_scores = res_data.get('rec_scores', [])

        # 确保三个列表长度一致
        min_length = min(len(dt_polys), len(rec_texts), len(rec_scores))

        for i in range(min_length):
            # 构造符合要求的JSON格式
            formatted_result = {
                "text": rec_texts[i],
                "bbox": dt_polys[i],
                "score": rec_scores[i]
            }
            data.append(formatted_result)
            # 输出格式化的JSON字符串
        logger.debug(json.dumps(data, ensure_ascii=False))
    return data
